package com.yociyy.security.store;

import com.yociyy.security.authentication.YoCiUser;
import com.yociyy.security.authentication.YoCiUserAuthenticationConverter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.provider.token.DefaultAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.JwtAccessTokenConverter;
import org.springframework.security.oauth2.provider.token.store.JwtTokenStore;
import org.springframework.security.oauth2.provider.token.store.KeyStoreKeyFactory;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * 自定义Oauth2 TokenStore配置类（采用JWT RSA 非对称加密算法）
 * 
 * @author: YoCiyy
 * @date: 2020/6/30
 */
@Slf4j
public class JWTTokenStore {

	@ConditionalOnProperty(prefix = "yoci.security.oauth2.tokenStore", name = "type", havingValue = "rsa-jwt-client")
	public static class RsaJwtClientTokenStore {

		@Bean
		public TokenStore tokenStore(JwtAccessTokenConverter jwtAccessTokenConverter) {
			log.info("tokenStore.type: rsa-jwt-client");
			return new JwtTokenStore(jwtAccessTokenConverter);
		}

		@Bean
		public JwtAccessTokenConverter jwtAccessTokenConverter() {

			JwtAccessTokenConverter accessTokenConverter = new JwtAccessTokenConverter();
			DefaultAccessTokenConverter tokenConverter = (DefaultAccessTokenConverter) accessTokenConverter.getAccessTokenConverter();
			tokenConverter.setUserTokenConverter(new YoCiUserAuthenticationConverter());

			Resource resource = new ClassPathResource("pubkey.txt");
			String pubKey;
			try (BufferedReader br = new BufferedReader(new InputStreamReader(resource.getInputStream()))) {
				pubKey = br.lines().collect(Collectors.joining("\n"));
			} catch (IOException e) {
				throw new RuntimeException();
			}
			accessTokenConverter.setVerifierKey(pubKey);
			return accessTokenConverter;
		}

	}

	@ConditionalOnProperty(prefix = "yoci.security.oauth2.tokenStore", name = "type", havingValue = "rsa-jwt-server")
	public static class RsaJwtServerTokenStore {

		@Bean
		public TokenStore tokenStore(JwtAccessTokenConverter jwtAccessTokenConverter) {
			log.info("tokenStore.type: rsa-jwt-server");
			return new JwtTokenStore(jwtAccessTokenConverter);
		}

		@Bean
		public JwtAccessTokenConverter jwtAccessTokenConverter() {
			JwtAccessTokenConverter accessTokenConverter = new JwtAccessTokenConverter();
			KeyStoreKeyFactory keyStoreKeyFactory = new KeyStoreKeyFactory(new ClassPathResource("yoci.jks"), "yoci-cloud".toCharArray());
			accessTokenConverter.setKeyPair(keyStoreKeyFactory.getKeyPair("yoci-cloud"));
			return accessTokenConverter;
		}

		/**
		 * jwt 生成token 定制化处理
		 * 添加一些额外的用户信息到token里面
		 */
		@Bean
		public TokenEnhancer tokenEnhancer() {
			return (accessToken, authentication) -> {
				final Map<String, Object> additionalInfo = new HashMap<>(1);
				Object principal = authentication.getUserAuthentication().getPrincipal();
				// 增加id参数
				if (principal instanceof YoCiUser) {
					YoCiUser user = (YoCiUser)principal;
					additionalInfo.put("id", user.getId());
					additionalInfo.put("deptId", user.getDeptId());
					additionalInfo.put("permissions", user.getPermissions());
				}
				((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(additionalInfo);
				return accessToken;
			};
		}
	}

}
